# Copyright (C) 2015-2024, Wazuh Inc.
# Created by Wazuh, Inc. <info@wazuh.com>.
# This program is free software; you can redistribute it and/or modify it under the terms of GPLv2
import os
import pytest
from pathlib import Path

from wazuh_testing.utils.db_queries import agent_db
from wazuh_testing.utils.file import read_json_file, copy, write_json_file, read_file, write_file
from wazuh_testing.modules.modulesd.vulnerability_detector import CUSTOM_CPE_HELPER_TEMPLATE
from wazuh_testing.constants.paths.databases import CPE_HELPER_PATH
from wazuh_testing import FEEDS_PATH


custom_cpe_helper_path = str(Path(FEEDS_PATH, 'cpe_helper', CUSTOM_CPE_HELPER_TEMPLATE))

tags_dictionary = {"VERSION_TAG": "version",
                   "FORMAT_TAG": "format_version",
                   "UPDATE_TAG": "update_date",
                   "DICTIONARY_TAG": "dictionary",
                   "TARGET_TAG": "target",
                   "SOURCE_TAG": "source",
                   "VENDOR_S_TAG": "vendor",
                   "PRODUCT_S_TAG": "product",
                   "VERSION_S_TAG": "version",
                   "TRANSLATION_TAG": "translation",
                   "VENDOR_T_TAG": "vendor",
                   "PRODUCT_T_TAG": "product",
                   "VERSION_T_TAG": "version",
                   "ACTION_TAG": "action",
                   "LICENSE_TAG": "license",
                   "TITLE_TAG": "title",
                   "COPYRIGHT_TAG": "copyright",
                   "DATE_TAG": "date",
                   "TYPE_TAG": "type"}

values_dictionary = {"VERSION_VALUE": "1.0",
                     "FORMAT_VALUE": "1.0",
                     "UPDATE_VALUE": "2050-10-02T10:56Z",
                     "TARGET_VALUE": "windows",
                     "VENDOR_S_VALUE": "^wazuh-mocking",
                     "PRODUCT_S_VALUE_0": "custom-package-0",
                     "PRODUCT_S_VALUE_1": "custom-package-1",
                     "VERSION_S_VALUE": "^custom-package-0",
                     "VENDOR_T_VALUE": "wazuh-mocking",
                     "PRODUCT_T_VALUE_0": "custom-package-0",
                     "PRODUCT_T_VALUE_1": "custom-package-1",
                     "VERSION_T_VALUE": "^custom-package.*",
                     "ACTION_VALUE_0": "replace_vendor",
                     "ACTION_VALUE_1": "replace_product",
                     "TITLE_VALUE": "Dictionary of CPEs to analyze system vulnerabilities.",
                     "COPYRIGHT_VALUE": "Copyright (C) 2015-2019, Wazuh Inc.",
                     "DATE_VALUE": "March 6, 2019.",
                     "TYPE_VALUE": "GPLv2"}


def replace_cpe_json_fields(test_metadata: dict[str, str], tags: dict[str, str] = None,
                            values: dict[str, str] = None) -> None:
    """Replace the tags and values of the generic_custom_cpe_helper.json file.

    Args:
        test_metadata (dict): Dictionary values to replace for specific test case
        tags (dict): Dictionary with tags names values
        values (dict): Dictionary with tag values
    """
    filedata = read_file(CPE_HELPER_PATH)
    for key in test_metadata['tags']:
        filedata = filedata.replace(key, test_metadata['tags'][key])
    for key in test_metadata['values']:
        filedata = filedata.replace(key, str(test_metadata['values'][key]))
    for key in tags:
        filedata = filedata.replace(key, str(tags[key]))

    for key in values:
        filedata = filedata.replace(key, str(values[key]))

    write_file(CPE_HELPER_PATH, data=filedata)


def remove_item(item: dict, remove_key: str = "") -> dict[str, str]:
    """Remove recursively the tags and values of the dictionary.

    Args:
        item (dict): Dictionary
        remove_key (str): Item to be deleted
    """
    if isinstance(item, dict):
        for key in list(item.keys()):
            if key == remove_key:
                del item[key]
            else:
                remove_item(item[key], remove_key)
    elif isinstance(item, list):
        for i in reversed(range(len(item))):
            if item[i] == remove_key:
                del item[i]
            else:
                remove_item(item[i], remove_key)
    else:
        pass
    return item


def remove_cpe_json_fields(tags: list = None) -> None:
    """Remove the tags and values of the generic_custom_cpe_helper.json file.

    Args:
        tags (list): Dictionary with tags names values
    """
    tags = [] if tags is None else tags

    if len(tags) > 0:
        filedata = read_json_file(CPE_HELPER_PATH)

        for item in tags:
            filedata = remove_item(filedata, item)

        write_json_file(CPE_HELPER_PATH, filedata)


@pytest.fixture()
def prepare_environment(test_metadata: dict[str, str], agent_system: str, agent_packages: list,
                        mock_agent_with_custom_system) -> int:
    """Prepare the environment with a mocked agent, vulnerable packages and a custom cpe_helper.

    - Mock an agent with a specified system.
    - Insert mocked vulnerables packages.
    - Update packages sync status.
    - Copy the custom CPE helper to the dictionaries folder.
    - Force full scan.

    Args:
        test_metadata (dict): Test case metadata.
        agent_system (str): System to set to the mocked agent.
        agent_packages (list): List of vulnerable packages
        mock_agent_with_custom_system (fixture): Mock an agent with a custom system.
    """
    missing_field = None if 'missing_field' not in test_metadata else test_metadata['missing_field']
    tags = None if 'tags' not in test_metadata else test_metadata['tags']
    values = None if 'values' not in test_metadata else test_metadata['values']

    for package in agent_packages:
        try:
            version = package['version']
        except KeyError:
            version = ''
        try:
            format = package['format']
        except KeyError:
            format = 'rpm'
        try:
            architecture = package['architecture']
        except KeyError:
            architecture = 'x64'
        agent_db.insert_package(name=package['name'], format=format, architecture=architecture,
                                agent_id=mock_agent_with_custom_system, vendor=package['vendor'],
                                version=version, source=package['name'])

    # Sync packages info
    agent_db.update_sync_info(agent_id=mock_agent_with_custom_system, component="syscollector-packages")
    agent_db.update_sync_info(agent_id=mock_agent_with_custom_system, component="syscollector-hotfixes")

    # Make a backup data from inital CPE helper
    cpe_helper_backup_data = read_json_file(CPE_HELPER_PATH)

    # Set the custom CPE helper
    copy(custom_cpe_helper_path, CPE_HELPER_PATH)

    # Remove the values of the CPE helper
    remove_cpe_json_fields(missing_field)

    # Replace the values of the CPE helper
    replace_cpe_json_fields(test_metadata=test_metadata, tags=tags_dictionary, values=values_dictionary)

    yield mock_agent_with_custom_system

    # Restore the CPE helper backup data
    write_json_file(CPE_HELPER_PATH, cpe_helper_backup_data)
